const fs = require('fs');
const { createCanvas, loadImage, ImageData } = require('canvas');
const ort = require('onnxruntime-node');

async function convertImage(options) {
    const {
        inputPath,
        outputPath,
        styleMode = 'Enhance_Edge',
        pixelSize = 'Small',
        edgeDensity = 'Large'
    } = options;

    // 创建canvas
    const canvas1 = createCanvas(1, 1);
    const ctx1 = canvas1.getContext('2d');
    const canvas2 = createCanvas(1, 1);
    const ctx2 = canvas2.getContext('2d');

    // 加载图片
    const img = await loadImage(inputPath);
    
    // 调整图片尺寸
    const imgArea = img.width * img.height;
    const maxArea = 2000 * 1100;
    let [canvas1Width, canvas1Height] = [img.width, img.height];
    
    if (imgArea > maxArea) {
        const scaleFactor = Math.sqrt(maxArea / imgArea);
        canvas1Width *= scaleFactor;
        canvas1Height *= scaleFactor;
    }
    
    canvas1.width = parseInt(canvas1Width);
    canvas1.height = parseInt(canvas1Height);
    ctx1.drawImage(img, 0, 0, canvas1.width, canvas1.height);

    // 获取图像数据
    const imageData = ctx1.getImageData(0, 0, canvas1.width, canvas1.height);
    const inputDataF32 = new Float32Array(imageData.data);

    // 创建Tensor
    const inputDataTensor = new ort.Tensor('float32', inputDataF32, [canvas1.height, canvas1.width, 4]);
    
    // 边缘密度参数
    const edgeThreshValues = {
        Small: 210,
        Medium: 168,
        Large: 128,
        Extra_Large: 80
    };

    const inputParam = new ort.Tensor('int32', Int32Array.from([edgeThreshValues[edgeDensity]]), [1]);

    // 模型路径映射
    const modelPaths = {
        Enhance_Edge: {
            Extra_Small: './onnx/model_8bit_edge_enchance_pixel_size_6_v1.onnx',
            Small: './onnx/model_8bit_edge_enchance_pixel_size_10_v1.onnx',
            Medium: './onnx/model_8bit_edge_enchance_pixel_size_15_v1.onnx',
            Large: './onnx/model_8bit_edge_enchance_pixel_size_20_v1.onnx',
            Extra_Large: './onnx/model_8bit_edge_enchance_pixel_size_28_v1.onnx'
        },
        Isolate_Pixel: {
            Extra_Small: './onnx/model_8bit_isolate_pixel_pixel_size_6_v1.onnx',
            Small: './onnx/model_8bit_isolate_pixel_pixel_size_10_v1.onnx',
            Medium: './onnx/model_8bit_isolate_pixel_pixel_size_15_v1.onnx',
            Large: './onnx/model_8bit_isolate_pixel_pixel_size_20_v1.onnx',
            Extra_Large: './onnx/model_8bit_isolate_pixel_pixel_size_28_v1.onnx'
        },
        Raw_Pixel: {
            Extra_Small: './onnx/model_8bit_raw_pixel_pixel_size_6_v1.onnx',
            Small: './onnx/model_8bit_raw_pixel_pixel_size_10_v1.onnx',
            Medium: './onnx/model_8bit_raw_pixel_pixel_size_15_v1.onnx',
            Large: './onnx/model_8bit_raw_pixel_pixel_size_20_v1.onnx',
            Extra_Large: './onnx/model_8bit_raw_pixel_pixel_size_28_v1.onnx'
        }
    };

    // 加载模型
    const session = await ort.InferenceSession.create(modelPaths[styleMode][pixelSize]);
    
    // 运行推理
    const feeds = styleMode === 'Enhance_Edge' 
        ? { input1: inputDataTensor, param1: inputParam }
        : { input1: inputDataTensor };

    const results = await session.run(feeds);

    // 处理输出
    const [resultHeight, resultWidth] = results.output1.dims;
    canvas2.width = resultWidth;
    canvas2.height = resultHeight;

    const outputDataUint8 = new Uint8ClampedArray(results.output1.data);
    const newImgData = new ImageData(outputDataUint8, resultWidth, resultHeight);
    ctx2.putImageData(newImgData, 0, 0);

    // 保存结果
    const out = fs.createWriteStream(outputPath);
    const stream = canvas2.createPNGStream();
    stream.pipe(out);

    return new Promise((resolve) => {
        out.on('finish', () => {
            console.log(`转换完成，结果已保存至 ${outputPath}`);
            resolve();
        });
    });
}

// 使用示例
convertImage({
    inputPath: './input.png',    // 输入图片路径
    outputPath: './output.png', // 输出图片路径
    styleMode: 'Raw_Pixel',  // 可选参数
    pixelSize: 'Medium',         // 可选参数
    edgeDensity: 'Large'        // 可选参数
}).catch(console.error);